from grid_env_ideal_obs_repeat_task import *
from grid_agent import *
from checkpoint_utils import *
from maze_factory import *
from replay_config import *
import argparse
import json
import sys
import matplotlib.pyplot as plt
from sklearn.manifold import TSNE
import matplotlib
import matplotlib.pyplot as plt
from matplotlib.patches import Circle
from matplotlib.lines import Line2D
from sklearn.manifold import TSNE
import random
from sklearn.decomposition import PCA
from matplotlib.animation import FuncAnimation
from sklearn.cluster import KMeans
import threading
import mplcursors
from mpl_toolkits.mplot3d.art3d import Poly3DCollection
from scipy.spatial.distance import pdist, squareform
from scipy.stats import pearsonr

def preprocess(trjs, max_length):
    processed_trjs = []
    for i in range(trjs.shape[0]):
        trj = trjs[i]
        if trj.shape[0] < max_length:
            last_element = trj[-1]
            processed_trj = np.concatenate([trj, np.repeat(last_element[np.newaxis,:], max_length - trj.shape[0], axis=0)], axis=0)
            processed_trjs.append(processed_trj)
        else:
            processed_trjs.append(trj)
    processed_trjs = np.array(processed_trjs)
    return processed_trjs

def main():

    pearsonrs_raw = np.load("./logs/pearsonrs_0.npy")
    pearsonrs_pca3 = np.load("./logs/pearsonrs_1.npy")
    capacities = np.load("./logs/capacities.npy")

    pearsonrs_diff = pearsonrs_raw - pearsonrs_pca3
    capacities = capacities/np.max(capacities)

    plt.plot(pearsonrs_raw, label="raw")
    plt.show()

    # 绘制这两条曲线
    plt.plot(pearsonrs_raw, label="raw")
    plt.plot(pearsonrs_pca3, label="pca3")
    plt.plot(pearsonrs_diff, label="diff")
    plt.plot(capacities, label="capacities")
    plt.legend()
    plt.show()

    """ parse arguments
    """
    rpl_config = ReplayConfig()

    parser = argparse.ArgumentParser()
    parser.add_argument("--model_pth", type=str, default=rpl_config.model_pth)
    parser.add_argument("--map_size", type=int, default=rpl_config.map_size)
    parser.add_argument("--task_pth", type=str, default=rpl_config.task_pth)
    parser.add_argument("--log_pth", type=str, default=rpl_config.log_pth)
    parser.add_argument("--nn_size", type=int, default=rpl_config.nn_size)
    parser.add_argument("--nn_type", type=str, default=rpl_config.nn_type)
    parser.add_argument("--show_kf", type=str, default=rpl_config.show_kf)
    parser.add_argument("--visualization", type=str, default=rpl_config.visualization)
    parser.add_argument("--video_output", type=str, default=rpl_config.video_output)
    parser.add_argument("--life_duration", type=int, default=rpl_config.life_duration)

    args = parser.parse_args()

    rpl_config.model_pth = args.model_pth
    rpl_config.map_size = args.map_size
    rpl_config.task_pth = args.task_pth
    rpl_config.log_pth = args.log_pth
    rpl_config.nn_size = args.nn_size
    rpl_config.nn_type = args.nn_type
    rpl_config.show_kf = args.show_kf
    rpl_config.visualization = args.visualization
    rpl_config.video_output = args.video_output
    rpl_config.life_duration = args.life_duration

    nn_type = ''
    if rpl_config.nn_type == "vanilla":
        nn_type = "vanilla"
    elif rpl_config.nn_type == "gru":
        nn_type = "gru"

    def load_data_and_compute(nn_type, seq_len, redundancy, diverse_set_capacity):

        rnn_limit_rings_file_name = "./logs/rnn_limit_rings_of_best_estimation_" + nn_type + "_" + str(seq_len) + "_" + str(redundancy) + "_" + str(diverse_set_capacity) + ".npz"
        # 载入 npz 文件
        rnn_limit_rings_of_best_estimation_file = np.load(rnn_limit_rings_file_name)

        # 获取 npz 文件中的所有对象名称
        matrix_names = rnn_limit_rings_of_best_estimation_file.files

        rnn_limit_rings_of_best_estimation = []

        # 遍历对象名称，访问和操作每个矩阵对象
        for name in matrix_names:
            matrix = rnn_limit_rings_of_best_estimation_file[name]
            # 在这里进行对矩阵对象的操作
            # 例如，打印矩阵的形状
            # print(f"Matrix '{name}' shape: {matrix.shape}")
            rnn_limit_rings_of_best_estimation.append(matrix)

        # 求 rnn_limit_rings_of_best_estimation 的中心位置序列
        rnn_limit_rings_of_best_estimation_center = []
        for i in range(len(rnn_limit_rings_of_best_estimation)):
            rnn_limit_rings_of_best_estimation_center.append(np.mean(rnn_limit_rings_of_best_estimation[i], axis=(0,1)))
        rnn_limit_rings_of_best_estimation_center = np.array(rnn_limit_rings_of_best_estimation_center)
        print("rnn_limit_rings_of_best_estimation_center.shape: ", rnn_limit_rings_of_best_estimation_center.shape)

        file_name = "obs_data_" + nn_type + "_" + str(seq_len) + "_" + str(redundancy) + "_" + str(diverse_set_capacity) + ".npz"
        obs_file = np.load("./logs/" + file_name)
        diverse_set_trajectoies = obs_file["diverse_set_trajectoies"]
        print("diverse_set_trajectoies.shape: ", diverse_set_trajectoies.shape)

        return rnn_limit_rings_of_best_estimation_center, diverse_set_trajectoies
    
    configs = [

        [nn_type, 6, 1, 100],
        [nn_type, 7, 1, 100],
        [nn_type, 8, 1, 100],
        [nn_type, 9, 1, 100],
        [nn_type, 10, 1, 100],
        [nn_type, 11, 1, 100],
        [nn_type, 12, 1, 100],
        [nn_type, 13, 1, 100],
        [nn_type, 14, 1, 100],
        [nn_type, 15, 1, 100],
        
        ]

    rnn_limit_rings_of_best_estimation_centers = []
    rnn_limit_rings_of_best_estimation_centers_0 = []
    ring_lengths = []
    diverse_set_trajectoies = []
    for i in range(len(configs)):
        _rnn_limit_rings_of_best_estimation_centers, dts = load_data_and_compute(configs[i][0], configs[i][1], configs[i][2], configs[i][3])
        rnn_limit_rings_of_best_estimation_centers.append(_rnn_limit_rings_of_best_estimation_centers)
        if i == 2:
            rnn_limit_rings_of_best_estimation_centers_0 = _rnn_limit_rings_of_best_estimation_centers.copy()
        diverse_set_trajectoies.append(dts)
        ring_length = configs[i][1]
        data_cap = _rnn_limit_rings_of_best_estimation_centers.shape[0]
        ring_length_ = np.array([ring_length for i in range(data_cap)])
        ring_lengths.append(ring_length_)

    ring_lengths = np.array(ring_lengths)
    ring_lengths = ring_lengths.reshape(ring_lengths.shape[0]*ring_lengths.shape[1])
    print(ring_lengths)

    diverse_set_trajectoies_mat = []
    for i in range(len(diverse_set_trajectoies)):
        length_aligned_trj = preprocess(diverse_set_trajectoies[i], 15)
        diverse_set_trajectoies_mat.append(length_aligned_trj)
        # print("diverse_set_trajectoies_mat[i].shape: ", diverse_set_trajectoies_mat[i].shape)
    diverse_set_trajectoies_mat = np.concatenate(diverse_set_trajectoies_mat, axis=0)
    print("diverse_set_trajectoies_mat.shape: ", diverse_set_trajectoies_mat.shape)


    # 将 rnn_limit_rings_of_best_estimation_centers 所有元素拼接起来
    rnn_limit_rings_of_best_estimation_center_mat = np.concatenate(rnn_limit_rings_of_best_estimation_centers, axis=0)
    print("rnn_limit_rings_of_best_estimation_center_mat.shape: ", rnn_limit_rings_of_best_estimation_center_mat.shape)

    # 对 rnn_limit_rings_of_best_estimation_center_mat 进行 PCA
    pca = PCA()
    pca.fit(rnn_limit_rings_of_best_estimation_center_mat)
    rnn_limit_rings_of_best_estimation_center_mat_pca = pca.transform(rnn_limit_rings_of_best_estimation_center_mat)
    print("rnn_limit_rings_of_best_estimation_center_mat_pca.shape: ", rnn_limit_rings_of_best_estimation_center_mat_pca.shape)

    pca0 = PCA()
    pca0.fit(rnn_limit_rings_of_best_estimation_centers_0)
    print("shapee of rnn_limit_rings_of_best_estimation_centers_0: ", rnn_limit_rings_of_best_estimation_centers_0.shape)

    # 计算角度映射
    rnn_limit_rings_of_best_estimation_center_mat_pca_2d = rnn_limit_rings_of_best_estimation_center_mat_pca[:, :2]
    # 计算 rnn_limit_rings_of_best_estimation_center_mat_pca_3d 中每个点相对于中心点的角度
    rnn_limit_rings_of_best_estimation_center_mat_pca_3d_angles = []
    for i in range(rnn_limit_rings_of_best_estimation_center_mat_pca_2d.shape[0]):
        rnn_limit_rings_of_best_estimation_center_mat_pca_3d_angles.append(np.arctan2(rnn_limit_rings_of_best_estimation_center_mat_pca_2d[i,1], rnn_limit_rings_of_best_estimation_center_mat_pca_2d[i,0]))
    rnn_limit_rings_of_best_estimation_center_mat_pca_3d_angles = np.array(rnn_limit_rings_of_best_estimation_center_mat_pca_3d_angles)
    # 将 rnn_limit_rings_of_best_estimation_center_mat_pca_3d_angles 转换到 [0, 2pi] 区间
    rnn_limit_rings_of_best_estimation_center_mat_pca_3d_angles = rnn_limit_rings_of_best_estimation_center_mat_pca_3d_angles - rnn_limit_rings_of_best_estimation_center_mat_pca_3d_angles.min()

    # 筛选出特定角度的数据
    """
        141
        189
        237
        223
        315
        135
        33
        289
    """

    # rnn_limit_rings_of_best_estimation_center_mat_pca = pca0.transform(rnn_limit_rings_of_best_estimation_center_mat)

    random_angle = 189
    random_angle *= (np.pi / 180)
    sector_angle = np.pi / 16
    # 选中所有角度在 [random_angle - sector_angle, random_angle + sector_angle] 之间的点
    idx = np.where((rnn_limit_rings_of_best_estimation_center_mat_pca_3d_angles >= random_angle - sector_angle) & (rnn_limit_rings_of_best_estimation_center_mat_pca_3d_angles <= random_angle + sector_angle))
    # print("idx", idx)
    rnn_limit_rings_of_best_estimation_center_mat_selected = rnn_limit_rings_of_best_estimation_center_mat_pca[idx]
    seq_len_selected = ring_lengths[idx]
    trj_selected = diverse_set_trajectoies_mat[idx]

    # 按照 ring_lengths 给 rnn_limit_rings_of_best_estimation_center_mat_pca 分配颜色值
    # 先创建一个长度为15的颜色列表
    colors = []
    # 指定随机种子为100
    np.random.seed(100)
    for i in range(15):
        colors.append(np.random.rand(3))
    
    # 给 rnn_limit_rings_of_best_estimation_center_mat_pca 中的每个点分配颜色
    rnn_limit_rings_of_best_estimation_center_mat_pca_colors = []
    for i in range(rnn_limit_rings_of_best_estimation_center_mat_selected.shape[0]):
        rnn_limit_rings_of_best_estimation_center_mat_pca_colors.append(colors[seq_len_selected[i]-1])

    # 计算 rnn_limit_rings_of_best_estimation_center_mat_selected 中具有相同的 seq_len_selected 的中心点
    rnn_limit_rings_of_best_estimation_center_mat_selected_centers = []
    for i in range(1, 16):
        idx = np.where(seq_len_selected == i)
        rnn_limit_rings_of_best_estimation_center_mat_selected_centers.append(np.mean(rnn_limit_rings_of_best_estimation_center_mat_selected[idx], axis=0))
        
    rnn_limit_rings_of_best_estimation_center_mat_selected_centers = np.array(rnn_limit_rings_of_best_estimation_center_mat_selected_centers)
    print("rnn_limit_rings_of_best_estimation_center_mat_selected_centers.shape: ", rnn_limit_rings_of_best_estimation_center_mat_selected_centers.shape)
    print(rnn_limit_rings_of_best_estimation_center_mat_selected_centers)

    # 使用不同的颜色将 rnn_limit_rings_of_best_estimation_centers_pca 中的所有数据进行可视化
    fig = plt.figure()
    ax = fig.add_subplot(111, projection='3d')
    
    vdim0 = 3
    vdim1 = 4
    vdim2 = 5

    ax.scatter(
        rnn_limit_rings_of_best_estimation_center_mat_selected[:,vdim0], 
        rnn_limit_rings_of_best_estimation_center_mat_selected[:,vdim1], 
        rnn_limit_rings_of_best_estimation_center_mat_selected[:,vdim2], 
        c=rnn_limit_rings_of_best_estimation_center_mat_pca_colors, 
        s=10,
        )
    
    ax.scatter(
        rnn_limit_rings_of_best_estimation_center_mat_selected_centers[:,vdim0], 
        rnn_limit_rings_of_best_estimation_center_mat_selected_centers[:,vdim1], 
        rnn_limit_rings_of_best_estimation_center_mat_selected_centers[:,vdim2], 
        c=colors, 
        s=500,
        marker='*',
        )
    
    print("shape of rnn_limit_rings_of_best_estimation_center_mat_selected_centers: ", rnn_limit_rings_of_best_estimation_center_mat_selected_centers.shape)
    
    # 在每个 rnn_limit_rings_of_best_estimation_center_mat_selected_centers 的位置，写上他们在 colors 中的索引
    for i in range(rnn_limit_rings_of_best_estimation_center_mat_selected_centers.shape[0]):
        ax.text(
            rnn_limit_rings_of_best_estimation_center_mat_selected_centers[i,vdim0], 
            rnn_limit_rings_of_best_estimation_center_mat_selected_centers[i,vdim1], 
            rnn_limit_rings_of_best_estimation_center_mat_selected_centers[i,vdim2], 
            str(i+1),
            fontsize=20,
            )
        
    ax.view_init(elev=90, azim=0)
    plt.show()

    # # 在2D平面上绘制 trj_selected
    
    # fig = plt.figure()
    # ax = fig.add_subplot(111)
    # for i in range(trj_selected.shape[0]-5, trj_selected.shape[0]-4):
    #     trj_selected[i] = trj_selected[i] - trj_selected[i][0]
    #     ax.plot(trj_selected[i,:,0], trj_selected[i,:,1], color=colors[seq_len_selected[i]-1], linewidth=2)
    # plt.show()

    
if __name__ == "__main__":
    main()